from path import Path
from pymongo import MongoClient
from tqdm import tqdm
from time import time, sleep
import re
import json
from multiprocessing import JoinableQueue, Process, current_process
import sys


def split(fp):
    with open(fp) as split_f:
        lines = []
        started = False

        for line in split_f:
            line = line.strip()
            if not started and line.startswith("<Entrezgene>"):
                started = True

            if started:
                lines.append(line)

            if started and (line.startswith("</Entrezgene>") or line.startswith("<Gene-commentary_products>")):
                yield ''.join(lines).replace("\n", '')
                lines.clear()
                started = False


def put_func(q, _dir):
    for _f in _dir.files("*.xml"):
        _class = _f.name[0:_f.name.rindex('_')]
        for _txt in split(_f):
            q.put(_class + "=xgxgx=" + _txt)


def get_func(q, iq):
    re_map = {
        "gene_track_geneid": ".*<Entrezgene_track-info>.*<Gene-track>.*<Gene-track_geneid>(.*)</Gene-track_geneid>",
        "org_ref_taxname": ".*<Entrezgene_source>.*<BioSource>.*<BioSource_org>.*<Org-ref>.*<Org-ref_taxname>(.*)</Org-ref_taxname>",
        "gene_ref_locus": ".*<Gene-ref>.*<Gene-ref_locus>(.*)</Gene-ref_locus>",
        "prot_ref_name_E": ".*<Entrezgene_prot>.*<Prot-ref>.*<Prot-ref_name>.*<Prot-ref_name_E>(.*)</Prot-ref_name_E>",
        "gene_source_src_int": ".*<Entrezgene_gene-source>.*<Gene-source>.*<Gene-source_src-int>(.*)</Gene-source_src-int>",
        "gene_source_src_str1": ".*<Entrezgene_gene-source>.*<Gene-source>.*<Gene-source_src-str1>(.*)</Gene-source_src-str1>",
        "gene_commentary_accession": ".*<Entrezgene_locus>.*<Gene-commentary>.*<Gene-commentary_accession>(.*)</Gene-commentary_accession>",
        "gene_commentary_version": ".*<Entrezgene_locus>.*<Gene-commentary>.*<Gene-commentary_version>(.*)</Gene-commentary_version>",
        "seq_interval_from": ".*<Entrezgene_locus>.*<Gene-commentary_seqs>.*<Seq-interval_from>(.*)</Seq-interval_from>",
        "seq_interval_to": ".*<Entrezgene_locus>.*<Gene-commentary_seqs>.*<Seq-interval_to>(.*)</Seq-interval_to>",
        "na_strand": '.*<Entrezgene_locus>.*<Gene-commentary_seqs>.*<Seq-interval_strand>.*<Na-strand value="(.*)"/>',
        "seq_id_gi": ".*<Entrezgene_locus>.*<Gene-commentary>.*<Gene-commentary_seqs>.*<Seq-loc>.*<Seq-loc_int>.*<Seq-interval>.*<Seq-interval_id>.*<Seq-id>.*<Seq-id_gi>(.*)</Seq-id_gi>",
    }

    def parse_entryz_gene(txt):
        txt = txt.split("=xgxgx=", 1)
        _class = txt[0]
        txt = txt[1]

        insert_data = {"class": _class}

        for k, v in re_map.items():
            res = re.match(v, txt)
            if res is not None and res.lastindex > 0:
                insert_data[k] = res.group(1)

        if len(insert_data):
            # 删掉空字符串
            for i in insert_data.keys():
                if not len(insert_data[i]):
                    del insert_data[i]

            iq.put(insert_data)

    while True:
        try:
            value = q.get()
            parse_entryz_gene(value)
        finally:
            q.task_done()


def insert_func(iq):
    client = MongoClient("localhost")
    tb = client["Pre_Data"]["PlantsGene"]

    while True:
        try:
            value = iq.get()
            tb.insert(value)
        finally:
            iq.task_done()


if __name__ == '__main__':
    start_time = time()

    process_size = 10
    dataDir = Path(sys.argv[1])
    if dataDir.exists():
        queue = JoinableQueue()
        insert_queue = JoinableQueue()

        # 文件拆分进程
        put_p = Process(target=put_func, args=(queue, dataDir))
        put_p.daemon = True
        put_p.start()

        # 文件解析进程
        for _i in range(process_size):
            gp = Process(target=get_func, args=(queue, insert_queue))
            gp.daemon = True
            gp.start()

        # mongodb 插入进程
        i_p = Process(target=insert_func, args=(insert_queue,))
        i_p.daemon = True
        i_p.start()

        while put_p.is_alive():
            log = "耗时 %ds | 待解析 %d | 待入库 %d" % (time() - start_time, queue.qsize(), insert_queue.qsize())
            sys.stdout.write("\r" + log)
            sys.stdout.flush()
            sleep(.5)

        print("\n拆分结束，等待解析和入库结束...")

        try:
            queue.join()
            insert_queue.join()
        except KeyboardInterrupt:
            print("stopped by hand")

    print("结束，程序耗费时间 %ds" % (time() - start_time))
